# -*- coding: utf-8 -*-  
'''
原始语料库

Created on 2021年9月8日
@author: luoyi
'''
import os
import random
import tqdm
import jieba_fast
import pickle

import utils.conf as conf
from utils.iexicon import LiteWordsWarehouse, WordsWarehouse, Punctuation, StopWords


#    原始文件数据集
class OriginalDataset:
    '''原始文件遍历'''
    def __init__(self,
                 dir_path=conf.DATASET_SOHU_THUCNEWS.get_original_dir_path(),
                 count=-1):
        '''
            @param dir_path: 原始文件目录
            @param count: 每个主题下遍历多少个文件，-1表示不限制
        '''
        self._dir_path = dir_path
        self._count = count
        pass
    
    #    所有文件迭代器
    def files_iterator(self, is_log=False):
        #    取所有一级子目录
        for topic_dir in os.listdir(self._dir_path):
            dpath = self._dir_path + '/' + topic_dir
            if (not os.path.isdir(dpath)): continue
            
            #    遍历主题目录下的文件
            files = os.listdir(dpath)
            fi = 0
            for file in files:
                fpath = dpath + '/' + file
                if (not os.path.isfile(fpath)): continue
                fi += 1
                if (self._count > 0 and fi > self._count): break
                
                yield topic_dir, fpath
                pass
            pass
        pass
    #    主题随机乱序遍历
    def files_iterator_shuffle(self):
        t_dirs = []
        t_dirs_idx = []                         #    主题目录的索引列表
        t_files = []                            #    每个主题目录下的文件列表
        t_files_idx = []                        #    每个主题目录下当前读到的文件索引
        t_files_len = []                        #    每个主题目录下的文件数
        i = 0
        for topic_dir in os.listdir(self._dir_path):
            dpath = self._dir_path + '/' + topic_dir
            if (not os.path.isdir(dpath)): continue
            
            t_dirs.append(topic_dir)            #    所有主题目录
            t_dirs_idx.append(i)                #    所有主题目录的索引
            i += 1
            files = os.listdir(dpath)           #    只遍历txt文件
            t_files.append(files)               #    每个主题下的所有文件
            t_files_len.append(len(files))      #    每个主题目录下文件数
            t_files_idx.append(0)               #    每个主题当前遍历指针
            pass
        
        #    开始遍历，条件是索引列表不为空。因为如果某个主题遍历完了会从列表中删除他的索引
        while (len(t_dirs_idx) > 0):
            #    随机取一个主题索引
            t_idx = random.randint(a=0, b=len(t_dirs_idx)-1)
            topic = t_dirs[t_idx]
            
            #    根据主题索引从他的文件列表中顺序取文件
            f_idx = t_files_idx[t_idx]
            file = t_files[t_idx][f_idx]
            
            #    更新指针
            f_idx += 1
            t_files_idx[t_idx] = f_idx
            #    判断是否已经读到末尾
            if ((f_idx >= t_files_len[t_idx])
                    or (self._count > 0 and f_idx >= self._count)): 
                del t_dirs_idx[t_idx]
                del t_dirs[t_idx]
                del t_files[t_idx]
                pass
            
            yield topic, self._dir_path + '/' + topic + '/' + file
            pass
        pass
    #    遍历每个主题下的每个文件
    def news_iterator(self, is_log=False):
        for topic, fpath in self.files_iterator(is_log):
            #    完整读一篇文档，读成句子格式
            sens = self.read_file(fpath)
            yield topic, sens
            pass
        pass
    #    遍历每个主题下的每个文件
    def news_iterator_shuffle(self):
        for topic, fpath in self.files_iterator_shuffle():
            #    完整读一篇文档，读成句子格式
            sens = self.read_file(fpath)
            yield topic, sens
            pass
        pass
    #    完整读取一篇文章
    def read_file(self, fpath, max_empty_line=10):
        '''完整读取一篇文章，以句子为单位。
            @return [句子1, 句子2...]
        '''
        doc_sens = []
        with open(file=fpath, mode='r', encoding='utf-8') as fr:
            line = fr.readline()
            empty_line = 0
            while (line):
                line = line.replace('\t', '').replace('\n', '').strip()

                if (len(line) == 0):
                    empty_line += 1
                    if (empty_line > max_empty_line): break
                    pass
                else:
                    empty_line = 0
                    
                    #    用标点符号断句，并去掉标点符号
                    sens = [sen for sen in Punctuation.split_by_sentence(line)]
                    doc_sens = doc_sens + sens
                    pass
                
                line = fr.readline()
                pass
            pass
        return doc_sens
    
    #    统计总文件数
    def stat_files(self):
        '''@return 主题数，每个主题的文章数，总文章数'''
        topic_count = 0
        topic_docs_count = []
        doc_count = 0
        for topic_dir in os.listdir(self._dir_path):
            dpath = self._dir_path + '/' + topic_dir
            if (not os.path.isdir(dpath)): continue
            
            topic_count += 1
            
            files = os.listdir(dpath)
            num_files = len(files)
            if (self._count > 0 and num_files > self._count): num_files = self._count
            
            doc_count += num_files
            topic_docs_count.append(num_files)
            pass
        
        return topic_count, topic_docs_count, doc_count
    
    #    统计词库
    def stat_words_warehouse(self, 
                             word_frequency_path=conf.DATASET_SOHU_THUCNEWS.get_word_frequency_path(),
                             word_id_path=conf.DATASET_SOHU_THUCNEWS.get_word_id_path(),
                             min_frequency=10):
        '''min_frequency：最小词频，超过最小词频的词会被过滤掉'''
        _, _, dc = self.stat_files()
        
        words_frequency = {}                                                 #    {k:词, v:词频}
        word_wid_dict = {}                                                   #    {k:词, v:词id}
        word_wid_dict.update(WordsWarehouse.default_word_wid())
        
        words_frequency_lite = {}
        word_wid_dict_lite = {}                                              #    {k:词, v:词id}，精简版 只保留词频超过阈值的词
        word_wid_dict_lite.update(WordsWarehouse.default_word_wid())
        
        #    先统计每个词的词频
        print('开始统计词频。。。')
        for _, sens in tqdm.tqdm(self.news_iterator(), total=dc, desc='统计文章', leave=True, ncols=100):
            #    对sens中每句话进行分词
            for s in sens:
                words = jieba_fast.lcut(s)                              #    完整分词
                words_lite = StopWords.instance().filter(words)         #    去掉停用词
                for word in words:
                    if (words_frequency.get(word) is None): words_frequency[word] = 1
                    else: words_frequency[word] = words_frequency[word] + 1
                    pass
                for word in words_lite:
                    if (words_frequency_lite.get(word) is None): words_frequency_lite[word] = 1
                    else: words_frequency_lite[word] = words_frequency_lite[word] + 1
                    pass
                pass
            pass
        print('词频统计完成。总词数：', len(words_frequency), ' 去掉停用词词数：', len(words_frequency_lite))
        
        
        print('\n开始生成各种文件。。。')
        
        print('写入词频文件。')
        #    降序排列
        l = sorted(words_frequency.items(), key=lambda item:item[1], reverse=True)
        #    写入词频文件
        fpath = word_frequency_path + '/' + WordsWarehouse.FNAME_WORD_FREQUENCY_TXT
        conf.mkfiledir_ifnot_exises(fpath)
        with open(file=fpath, mode='w', encoding='utf-8') as fw:
            for item in l:
                word = item[0]              #    词
                frequency = item[1]         #    频
                
                if (frequency <= min_frequency): continue
                
                fw.write('{}:{}\n'.format(word, frequency))
                
                #    过滤掉词频小于阈值的词
                word_wid_dict[word] = len(word_wid_dict)
                pass
            pass
        wid_word_dict = {v:k for k,v in word_wid_dict.items()}
        print('写入词库文件。。。')
        conf.mkfiledir_ifnot_exises(word_id_path + '/' + WordsWarehouse.FNAME_WID_WORD)
        with open(file=word_id_path + '/' + WordsWarehouse.FNAME_WID_WORD, mode='wb') as fw: pickle.dump(wid_word_dict, fw)
        conf.mkfiledir_ifnot_exises(word_id_path + '/' + WordsWarehouse.FNAME_WORD_WID)
        with open(file=word_id_path + '/' + WordsWarehouse.FNAME_WORD_WID, mode='wb') as fw: pickle.dump(word_wid_dict, fw)
        
        
        print('写入精简词频文件。')
        #    降序排列
        l = sorted(words_frequency_lite.items(), key=lambda item:item[1], reverse=True)
        #    写入词频文件
        fpath = word_frequency_path + '/' + LiteWordsWarehouse.FNAME_WORD_FREQUENCY_LITE_TXT
        conf.mkfiledir_ifnot_exises(fpath)
        with open(file=fpath, mode='w', encoding='utf-8') as fw:
            for item in l:
                word = item[0]              #    词
                frequency = item[1]         #    频
                
                if (frequency <= min_frequency): continue
                
                fw.write('{}:{}\n'.format(word, frequency))
                
                #    过滤掉词频小于阈值的词
                word_wid_dict_lite[word] = len(word_wid_dict_lite)
                pass
            pass
        wid_word_dict_lite = {v:k for k,v in word_wid_dict_lite.items()}
        print('写入精简词库文件。。。')
        conf.mkfiledir_ifnot_exises(word_id_path + '/' + LiteWordsWarehouse.FNAME_WID_WORD_LITE)
        with open(file=word_id_path + '/' + LiteWordsWarehouse.FNAME_WID_WORD_LITE, mode='wb') as fw: pickle.dump(wid_word_dict_lite, fw)
        conf.mkfiledir_ifnot_exises(word_id_path + '/' + LiteWordsWarehouse.FNAME_WORD_WID_LITE)
        with open(file=word_id_path + '/' + LiteWordsWarehouse.FNAME_WORD_WID_LITE, mode='wb') as fw: pickle.dump(word_wid_dict_lite, fw)
        
        print('词库文件写入完成。')
        print('完整词库（去掉低频词）：', len(wid_word_dict))
        print('精简词库（去掉停用词，低频词）：', len(word_wid_dict_lite))
        pass
    
    pass